﻿using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Drawing;
using System.Text;
using CSJ2K.Util;
using CSJ2K.Util.IO;

namespace CSJ2K.Jpc
{
    public class TagTree
    {
        uint[][,] Tree;
        bool[][,] Fixed;

        public TagTree(uint[,] Data)
            : this(Data.GetLength(0), Data.GetLength(1))
        {
            Tree[Tree.Length - 1] = Data;

            // Calculate values for lower layers
            for (int i = Tree.Length - 2; i >= 0; i--)
            {
                for (int x = 0; x < Tree[i].GetLength(0); x++)
                {
                    for (int y = 0; y < Tree[i].GetLength(1); y++)
                    {
                        uint min = Tree[i + 1][x * 2, y * 2];
                        if ((x * 2 + 1) < Tree[i + 1].GetLength(0))
                        {
                            min = Math.Min(min, Tree[i + 1][x * 2, y * 2]);
                            if ((y * 2 + 1) < Tree[i + 1].GetLength(1))
                                min = Math.Min(min, Math.Min(Tree[i + 1][x * 2, y * 2 + 1], Tree[i + 1][x * 2 + 1, y * 2 + 1]));
                        }
                        else if ((y * 2 + 1) < Tree[i + 1].GetLength(1)) min = Math.Min(min, Tree[i + 1][x * 2, y * 2 + 1]);

                        Tree[i][x, y] = min;
                    }
                }
            }

            // Allocate space for Fixed flags to record what has already been encoded.
            Fixed = new bool[Tree.Length][,];
            for (int i = 0; i < Tree.Length; i++) Fixed[i] = new bool[Tree[i].GetLength(0), Tree[i].GetLength(1)];
        }

        public TagTree(int X, int Y)
        {
            // Calculate the number of layers required based on the maximum dimension
            int Depth = 1;
            for (int i = Math.Max(X, Y); i > 1; i = (i + 1) / 2) Depth++;

            // Allocate space for each layer
            Tree = new uint[Depth][,];
            Tree[Depth - 1] = new uint[X, Y];
            for (int i = Depth - 2; i >= 0; i--)
            {
                int x = (Tree[i + 1].GetLength(0) + 1) / 2;
                int y = (Tree[i + 1].GetLength(1) + 1) / 2;
                Tree[i] = new uint[x, y];
            }

            // Allocate space for Fixed flags to record what has already been encoded.
            Fixed = new bool[Tree.Length][,];
            for (int i = 0; i < Tree.Length; i++) Fixed[i] = new bool[Tree[i].GetLength(0), Tree[i].GetLength(1)];

        }
        // TODO: support BitstreamWriter when it's writen.
        public void WriteValue(int X, int Y)
        {
            Point[] path = GetPath(X, Y);
            uint lastvalue = 0;
            for (int i = 0; i < path.Length; i++)
            {
                if (!Fixed[i][path[i].X, path[i].Y])
                {
                    uint value = Tree[i][path[i].X, path[i].Y] - lastvalue;
                    for (int j = 0; j < value; j++) Console.Write(0);
                    Console.Write(1);

                    Fixed[i][path[i].X, path[i].Y] = true;
                }
                lastvalue = Tree[i][path[i].X, path[i].Y];
            }
        }
        public void WriteIncluded(int X, int Y, int Layer)
        {
            Point[] path = GetPath(X, Y);
            uint lastvalue = 0;
            for (int i = 0; i < path.Length; i++)
            {
                if (!Fixed[i][path[i].X, path[i].Y])
                {
                    uint value = Tree[i][path[i].X, path[i].Y] - lastvalue;
                    for (int j = 0; j < value; j++) Console.Write(0);
                    Console.Write(1);

                    Fixed[i][path[i].X, path[i].Y] = true;

                    if (value > Layer) return;
                }
                lastvalue = Tree[i][path[i].X, path[i].Y];
            }
        }

        public override string ToString()
        {
            StringBuilder sb = new StringBuilder();
            for (int i = Tree.Length - 1; i >= 0; i--)
            {
                sb.AppendLine("Layer " + i);
                for (int y = 0; y < Tree[i].GetLength(1); y++)
                {
                    for (int x = 0; x < Tree[i].GetLength(0); x++)
                    {
                        if (x != 0) sb.Append(", ");
                        else sb.Append("\t");
                        if (Fixed[i][x, y]) sb.Append(Tree[i][x, y]);
                        else sb.Append("-");
                    }
                    sb.AppendLine();
                }
            }
            return sb.ToString();
        }
        public uint ReadValue(BitstreamReader br, int X, int Y)
        {
            uint value = 0;
            Point[] path = GetPath(X, Y);
            if (DEBUG.TagTree)
            {
                Debug.Write("Level Path: ");
                for (int i = 0; i < path.Length; i++) Debug.Write("[" + path[i].X + "," + path[i].Y + "]");
                Debug.WriteLine("");
            }
            for (int i = 0; i < path.Length; i++)
            {
                if (!Fixed[i][path[i].X, path[i].Y])
                {
                    if (DEBUG.TagTree) Debug.Write("READ bits for level " + i + " [" + path[i].X + "," + path[i].Y + "]: ");
                    while (!br.ReadBit())
                    {
                        if (DEBUG.TagTree) Debug.Write(0);
                        value++;
                    }
                    if (DEBUG.TagTree) Debug.Write(1);
                    Tree[i][path[i].X, path[i].Y] = value;
                    if (DEBUG.TagTree) Debug.WriteLine(": " + value);

                    Fixed[i][path[i].X, path[i].Y] = true;
                }
                else
                {
                    value = Tree[i][path[i].X, path[i].Y];
                    if (DEBUG.TagTree) Debug.WriteLine("Have bits for level " + i + " [" + path[i].X + "," + path[i].Y + "]: " + value);
                }
            }
            return value;
        }

        // Special truncated tag tree reader implementation for inclusion tag-trees
        public bool ReadIncluded(BitstreamReader br, int X, int Y, int Layer)
        {
            //Debug.Write(this.ToString());

            uint value = 0;
            Point[] path = GetPath(X, Y);
            if (DEBUG.TagTree)
            {
                Debug.Write("Level Path: ");
                for (int i = 0; i < path.Length; i++) Debug.Write("[" + path[i].X + "," + path[i].Y + "]");
                Debug.WriteLine("");
            }
            for (int i = 0; i < path.Length; i++)
            {
                if (!Fixed[i][path[i].X, path[i].Y])
                {
                    bool truncated = false;

                    //value = Tree[i][path[i].X, path[i].Y];
                    if (Tree[i][path[i].X, path[i].Y] > 0)
                    {
                        value = Tree[i][path[i].X, path[i].Y];
                        if (DEBUG.TagTree) Debug.WriteLine("Previously truncated. (" + value + ")");
                        if (value > Layer)
                        {
                            if (DEBUG.TagTree) Debug.WriteLine("Trucating again.");
                            truncated = true;
                        }
                    }

                    if (!truncated)
                    {
                        if (DEBUG.TagTree) Debug.Write("READ bits for level " + i + " [" + path[i].X + "," + path[i].Y + "]: ");
                        while (!truncated && !br.ReadBit())
                        {
                            if (DEBUG.TagTree) Debug.Write(0);
                            value++;
                            if (value > Layer) truncated = true;
                        }

                        Tree[i][path[i].X, path[i].Y] = value;

                        if (!truncated)
                        {
                            if (DEBUG.TagTree) Debug.Write(1);
                            Fixed[i][path[i].X, path[i].Y] = true;
                            if (DEBUG.TagTree) Debug.WriteLine(" : " + value);
                        }
                        else
                        {
                            if (DEBUG.TagTree) Debug.WriteLine(" : >=" + value + " (Truncated)");
                        }
                    }
                }
                else
                {
                    value = Tree[i][path[i].X, path[i].Y];
                    if (DEBUG.TagTree) Debug.WriteLine("Have bits for level " + i + " [" + path[i].X + "," + path[i].Y + "]: " + value);

                    // if this is the last layer and we already have bits, read a single 0 or 1 for included.
                    if (i == path.Length - 1)
                    {
                        if (DEBUG.TagTree) Debug.WriteLine("Reading inclusion bit.");
                        if (br.ReadBit()) return true;
                        else return false;
                    }
                }
                // Writing would be truncated here if we already know the block isn't included yet.
                if (value > Layer)
                {
                    if (DEBUG.TagTree) Debug.WriteLine("Already know it's not included.");
                    return false;
                }
            }
            if (value > Layer) return false;
            else return true;
        }

        private Point[] GetPath(int X, int Y)
        {
            Point[] path = new Point[Tree.Length];

            path[Tree.Length - 1].X = X;
            path[Tree.Length - 1].Y = Y;
            for (int i = Tree.Length - 2; i >= 0; i--)
            {
                path[i].X = path[i + 1].X / 2;
                path[i].Y = path[i + 1].Y / 2;
            }
            return path;
        }

    }
}
